package ollama

import (
	"context"
	"net/url"
	"strings"
	"time"

	"gitee.com/linxing_3/sr-core/plugins/storage/cache/file"
	"github.com/ollama/ollama/envconfig"
	"github.com/spf13/cast"
)

const ollamaEnvConfigKey = "ollama_config"

type EnvCfg struct {
	OLLAMA_DEBUG             bool
	OLLAMA_FLASH_ATTENTION   bool
	OLLAMA_KV_CACHE_TYPE     string
	OLLAMA_GPU_OVERHEAD      uint64
	OLLAMA_HOST              *url.URL
	OLLAMA_KEEP_ALIVE        time.Duration
	OLLAMA_LLM_LIBRARY       string
	OLLAMA_LOAD_TIMEOUT      time.Duration
	OLLAMA_MAX_LOADED_MODELS uint
	OLLAMA_MAX_QUEUE         uint
	OLLAMA_MODELS            string
	OLLAMA_NOHISTORY         bool
	OLLAMA_NOPRUNE           bool
	OLLAMA_NUM_PARALLEL      uint
	OLLAMA_ORIGINS           []string
	OLLAMA_SCHED_SPREAD      bool
	OLLAMA_MULTIUSER_CACHE   bool
	OLLAMA_NEW_ENGINE        bool
	Informational            string
	HTTP_PROXY               string
	HTTPS_PROXY              string
	NO_PROXY                 string

	cache *file.File
}

func (e *EnvCfg) Init() {
	if maVal, err := e.cache.HashGet(context.Background(), ollamaEnvConfigKey); err == nil && len(maVal) > 0 {
		e.loadMap(maVal)
		return
	}
	e.OLLAMA_DEBUG = envconfig.Debug()
	e.OLLAMA_FLASH_ATTENTION = envconfig.FlashAttention()
	e.OLLAMA_KV_CACHE_TYPE = envconfig.KvCacheType()
	e.OLLAMA_GPU_OVERHEAD = envconfig.GpuOverhead()
	e.OLLAMA_HOST = envconfig.Host()
	e.OLLAMA_KEEP_ALIVE = envconfig.KeepAlive()
	e.OLLAMA_LLM_LIBRARY = envconfig.LLMLibrary()
	e.OLLAMA_LOAD_TIMEOUT = envconfig.LoadTimeout()
	e.OLLAMA_MAX_LOADED_MODELS = envconfig.MaxRunners()
	e.OLLAMA_MAX_QUEUE = envconfig.MaxQueue()
	e.OLLAMA_MODELS = envconfig.Models()
	e.OLLAMA_NOHISTORY = envconfig.NoHistory()
	e.OLLAMA_NOPRUNE = envconfig.NoPrune()
	e.OLLAMA_NUM_PARALLEL = envconfig.NumParallel()
	e.OLLAMA_ORIGINS = envconfig.Origins()
	e.OLLAMA_SCHED_SPREAD = envconfig.SchedSpread()
	e.OLLAMA_MULTIUSER_CACHE = envconfig.MultiUserCache()
	e.OLLAMA_NEW_ENGINE = envconfig.NewEngine()
	e.HTTP_PROXY = envconfig.String("HTTP_PROXY")()
	e.HTTPS_PROXY = envconfig.String("HTTPS_PROXY")()
	e.NO_PROXY = envconfig.String("NO_PROXY")()
}

func (e *EnvCfg) asMap() map[string]any {
	return map[string]any{
		"OLLAMA_DEBUG":             cast.ToString(e.OLLAMA_DEBUG),
		"OLLAMA_FLASH_ATTENTION":   cast.ToString(e.OLLAMA_FLASH_ATTENTION),
		"OLLAMA_KV_CACHE_TYPE":     cast.ToString(e.OLLAMA_KV_CACHE_TYPE),
		"OLLAMA_GPU_OVERHEAD":      cast.ToString(e.OLLAMA_GPU_OVERHEAD),
		"OLLAMA_HOST":              cast.ToString(e.OLLAMA_HOST),
		"OLLAMA_KEEP_ALIVE":        cast.ToString(e.OLLAMA_KEEP_ALIVE),
		"OLLAMA_LLM_LIBRARY":       cast.ToString(e.OLLAMA_LLM_LIBRARY),
		"OLLAMA_LOAD_TIMEOUT":      cast.ToString(e.OLLAMA_LOAD_TIMEOUT),
		"OLLAMA_MAX_LOADED_MODELS": cast.ToString(e.OLLAMA_MAX_LOADED_MODELS),
		"OLLAMA_MAX_QUEUE":         cast.ToString(e.OLLAMA_MAX_QUEUE),
		"OLLAMA_MODELS":            cast.ToString(e.OLLAMA_MODELS),
		"OLLAMA_NOHISTORY":         cast.ToString(e.OLLAMA_NOHISTORY),
		"OLLAMA_NOPRUNE":           cast.ToString(e.OLLAMA_NOPRUNE),
		"OLLAMA_NUM_PARALLEL":      cast.ToString(e.OLLAMA_NUM_PARALLEL),
		"OLLAMA_ORIGINS":           strings.Join(e.OLLAMA_ORIGINS, ","),
		"OLLAMA_SCHED_SPREAD":      cast.ToString(e.OLLAMA_SCHED_SPREAD),
		"OLLAMA_MULTIUSER_CACHE":   cast.ToString(e.OLLAMA_MULTIUSER_CACHE),
		"OLLAMA_NEW_ENGINE":        cast.ToString(e.OLLAMA_NEW_ENGINE),
		"HTTP_PROXY":               cast.ToString(e.HTTP_PROXY),
		"HTTPS_PROXY":              cast.ToString(e.HTTPS_PROXY),
		"NO_PROXY":                 cast.ToString(e.NO_PROXY),
	}
}

func (e *EnvCfg) loadMap(m map[string]string) {
	e.OLLAMA_DEBUG = cast.ToBool(m["OLLAMA_DEBUG"])
	e.OLLAMA_FLASH_ATTENTION = cast.ToBool(m["OLLAMA_FLASH_ATTENTION"])
	e.OLLAMA_KV_CACHE_TYPE = cast.ToString(m["OLLAMA_KV_CACHE_TYPE"])
	e.OLLAMA_GPU_OVERHEAD = cast.ToUint64(m["OLLAMA_GPU_OVERHEAD"])
	e.OLLAMA_HOST, _ = url.Parse(m["OLLAMA_HOST"])
	e.OLLAMA_KEEP_ALIVE = cast.ToDuration(m["OLLAMA_KEEP_ALIVE"])
	e.OLLAMA_LLM_LIBRARY = cast.ToString(m["OLLAMA_LLM_LIBRARY"])
	e.OLLAMA_LOAD_TIMEOUT = cast.ToDuration(m["OLLAMA_LOAD_TIMEOUT"])
	e.OLLAMA_MAX_LOADED_MODELS = cast.ToUint(m["OLLAMA_MAX_LOADED_MODELS"])
	e.OLLAMA_MAX_QUEUE = cast.ToUint(m["OLLAMA_MAX_QUEUE"])
	e.OLLAMA_MODELS = cast.ToString(m["OLLAMA_MODELS"])
	e.OLLAMA_NOHISTORY = cast.ToBool(m["OLLAMA_NOHISTORY"])
	e.OLLAMA_NOPRUNE = cast.ToBool(m["OLLAMA_NOPRUNE"])
	e.OLLAMA_NUM_PARALLEL = cast.ToUint(m["OLLAMA_NUM_PARALLEL"])
	e.OLLAMA_ORIGINS = strings.Split(cast.ToString(m["OLLAMA_ORIGINS"]), ",")
	e.OLLAMA_SCHED_SPREAD = cast.ToBool(m["OLLAMA_SCHED_SPREAD"])
	e.OLLAMA_MULTIUSER_CACHE = cast.ToBool(m["OLLAMA_MULTIUSER_CACHE"])
	e.OLLAMA_NEW_ENGINE = cast.ToBool(m["OLLAMA_NEW_ENGINE"])
	e.HTTP_PROXY = cast.ToString(m["HTTP_PROXY"])
	e.HTTPS_PROXY = cast.ToString(m["HTTPS_PROXY"])
	e.NO_PROXY = cast.ToString(m["NO_PROXY"])
}

func (e *EnvCfg) Flush(ctx context.Context) error {
	return e.cache.HashSet(ctx, ollamaEnvConfigKey, e.asMap())
}
